#!/usr/bin/env python
# -*- coding: utf-8 -*-  
# Copyright (c) 2011 flykeysky<at>gmail.com. All rights reserved.
# This program or module is free software: you can redistribute it and/or
# modify it under the terms of the GNU General Public License as published
# by the Free Software Foundation, either version 2 of the License, or
# version 3 of the License, or (at your option) any later version. It is
# provided for educational purposes and is distributed in the hope that
# it will be useful, but WITHOUT ANY WARRANTY; without even the implied
# warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See
# the GNU General Public License for more details.

from __future__ import division
from __future__ import print_function
from __future__ import unicode_literals
from future_builtins import *

import platform
from PyQt4.QtCore import (QAbstractTableModel, QDataStream, QFile,
        QIODevice, QModelIndex, QRegExp, QSize, QString, QVariant, Qt,
        SIGNAL)
from PyQt4.QtGui import (QApplication, QColor, QComboBox, QLineEdit,
        QSpinBox, QStyle, QStyledItemDelegate, QTextDocument, QTextEdit)
import richtextlineedit


OPENDPI, L4, PROCESSNAME, FIVETUPLE, FEATURES, PAYLOADFILENAME = range(6)

MAGIC_NUMBER = 0x570C4
FILE_VERSION = 1


class Flow(object):

    def __init__(self, opendpi="unknown", l4="unknown_L4", processname="unknown_process", \
                 fivetuple="0.0.0.0", features="", payload=""):
        self.opendpi = QString(opendpi)
        self.l4 = QString(l4)
        self.processname = QString(processname)
        self.fivetuple = QString(fivetuple)
        self.features = QString(features)
        self.payloadfilename = QString(payload)


    def __hash__(self):
        return super(Flow, self).__hash__()


    def __lt__(self, other):
        r = QString.localeAwareCompare(self.fivetuple.toLower(),
                                       other.fivetuple.toLower())
        return True if r < 0 else False


    def __eq__(self, other):
        return 0 == QString.localeAwareCompare(self.fivetuple.toLower(),
                                               other.fivetuple.toLower())


class FlowContainer(object):

    def __init__(self, filename=QString()):
        self.filename = QString(filename)
        self.dirty = False
        self.flows = {}
        self.l4set = set()
        self.processnameset = set()
        self.featuresset = set()
        self.payloadfilenameset = set()


    def flow(self, identity):
        return self.flows.get(identity)

        
    def addFlow(self, flow):
        self.flows[id(flow)] = flow
        self.l4set.add(unicode(flow.l4))
        self.processnameset.add(unicode(flow.processname))
        self.featuresset.add(flow.features)
        self.payloadfilenameset.add(flow.payloadfilename)
        self.dirty = True


    def removeFlow(self, flow):
        del self.flows[id(flow)]
        del flow
        self.dirty = True
    
    def emptyFlows(self):
        self.flows = {}

    def __len__(self):
        return len(self.flows)


    def __iter__(self):
        for flow in self.flows.values():
            yield flow


    def inOrder(self):
        return sorted(self.flows.values())


    def inProcessL4Order(self):
        return sorted(self.flows.values(),
                      key=lambda x: (x.processname, x.l4, x.opendpi))


    def load(self):
        exception = None
        fh = None
        try:
            if self.filename.isEmpty():
                raise IOError, "no filename specified for loading"
            fh = QFile(self.filename)
            if not fh.open(QIODevice.ReadOnly):
                raise IOError, unicode(fh.errorString())
            stream = QDataStream(fh)
            magic = stream.readInt32()
            if magic != MAGIC_NUMBER:
                raise IOError, "unrecognized file type"
            fileVersion = stream.readInt16()
            if fileVersion != FILE_VERSION:
                raise IOError, "unrecognized file type version"
            self.flows = {}
            while not stream.atEnd():
                opendpi = QString()
                l4 = QString()
                processname = QString()
                fivetuple = QString()
                features =  QString()
                payload = QString()
                stream >> opendpi >> l4 >> processname >> fivetuple >> features >>payload
                #features = stream.readInt32()
                flow = Flow(opendpi, l4, processname, fivetuple, features, payload)
                self.flows[id(flow)] = flow
                self.l4set.add(unicode(l4))
                self.processnameset.add(unicode(processname))
                self.featuresset.add(features)
                self.payloadfilenameset.add(payload)
            self.dirty = False
        except IOError, e:
            exception = e
        finally:
            if fh is not None:
                fh.close()
            if exception is not None:
                raise exception


    def save(self):
        exception = None
        fh = None
        try:
            if self.filename.isEmpty():
                raise IOError, "no filename specified for saving"
            fh = QFile(self.filename)
            if not fh.open(QIODevice.WriteOnly):
                raise IOError, unicode(fh.errorString())
            stream = QDataStream(fh)
            stream.writeInt32(MAGIC_NUMBER)
            stream.writeInt16(FILE_VERSION)
            stream.setVersion(QDataStream.Qt_4_1)
            for flow in self.flows.values():
                stream << flow.opendpi << flow.l4 << flow.processname \
                       << flow.fivetuple << flow.features <<flow.payloadfilename
                #stream.writeInt32(flow.features)
            self.dirty = False
        except IOError, e:
            exception = e
        finally:
            if fh is not None:
                fh.close()
            if exception is not None:
                raise exception


class FlowTableModel(QAbstractTableModel):

    def __init__(self, filename=QString()):
        super(FlowTableModel, self).__init__()
        self.filename = filename
        self.dirty = False
        self.flows = []
        self.l4set = set()
        self.processnameset = set()
        self.featuresset = set()
        self.payloadfilenameset = set()


    def sortByName(self):
        self.flows = sorted(self.flows)
        self.reset()


    def sortByCountryOwner(self):
        self.flows = sorted(self.flows,
                            key=lambda x: (x.processname, x.l4, x.opendpi))
        self.reset()


    def flags(self, index):
        if not index.isValid():
            return Qt.ItemIsEnabled
        return Qt.ItemFlags(
                QAbstractTableModel.flags(self, index)|
                Qt.ItemIsEditable)


    def data(self, index, role=Qt.DisplayRole):
        if (not index.isValid() or
            not (0 <= index.row() < len(self.flows))):
            return QVariant()
        flow = self.flows[index.row()]
        column = index.column()
        if role == Qt.DisplayRole:
            if column == OPENDPI:
                return QVariant(flow.opendpi)
            elif column == L4:
                return QVariant(flow.l4)
            elif column == PROCESSNAME:
                return QVariant(flow.processname)
            elif column == FIVETUPLE:
                return QVariant(flow.fivetuple)
            elif column == FEATURES:
                return QVariant(flow.features )
            elif column == PAYLOADFILENAME:
                return QVariant(flow.payloadfilename )
        elif role == Qt.TextAlignmentRole:
            if column == FEATURES:
                return QVariant( int(Qt.AlignRight|Qt.AlignVCenter) )
            return QVariant( int(Qt.AlignLeft|Qt.AlignVCenter) )
        elif role == Qt.TextColorRole and column == FEATURES:
            return QVariant(QColor(Qt.red))
        elif role == Qt.BackgroundColorRole:
            if flow.processname in ("thunder.exe"):
                return QVariant(QColor(250, 230, 250))
            elif flow.processname in ("pplive.exe"):
                return QVariant(QColor(250, 250, 230))
            elif flow.processname in ("others",):
                return QVariant(QColor(230, 250, 250))
            else:
                return QVariant(QColor(210, 230, 230))
        return QVariant()


    def headerData(self, section, orientation, role=Qt.DisplayRole):
        if role == Qt.TextAlignmentRole:
            if orientation == Qt.Horizontal:
                return QVariant(int(Qt.AlignLeft|Qt.AlignVCenter))
            return QVariant(int(Qt.AlignRight|Qt.AlignVCenter))
        if role != Qt.DisplayRole:
            return QVariant()
        if orientation == Qt.Horizontal:
            if section == OPENDPI:
                return QVariant("OpenDPI")
            elif section == L4:
                return QVariant("L4")
            elif section == PROCESSNAME:
                return QVariant("ProcessName")
            elif section == FIVETUPLE:
                return QVariant("Fivetuple")
            elif section == FEATURES:
                return QVariant("Features")
            elif section == PAYLOADFILENAME:
                return QVariant("Payloadfilename")
        return QVariant(int(section + 1))


    def rowCount(self, index=QModelIndex()):
        return len(self.flows)


    def columnCount(self, index=QModelIndex()):
        return 5


    def setData(self, index, value, role=Qt.EditRole):
        if index.isValid() and 0 <= index.row() < len(self.flows):
            flow = self.flows[index.row()]
            column = index.column()
            if column == OPENDPI:
                flow.opendpi = value.toString()
            elif column == L4:
                flow.l4 = value.toString()
            elif column == PROCESSNAME:
                flow.processname = value.toString()
            elif column == FIVETUPLE:
                flow.fivetuple = value.toString()
            elif column == FEATURES:
                flow.features = value.toString()
            elif column == PAYLOADFILENAME:
                flow.payloadfilename = value.toString()
            self.dirty = True
            self.emit(SIGNAL("dataChanged(QModelIndex,QModelIndex)"),
                      index, index)
            return True
        return False


    def insertRows(self, position, rows=1, index=QModelIndex()):
        self.beginInsertRows(QModelIndex(), position, position + rows - 1)
        for row in range(rows):
            self.flows.insert(position + row,
                              Flow())
        self.endInsertRows()
        self.dirty = True
        return True


    def removeRows(self, position, rows=1, index=QModelIndex()):
        self.beginRemoveRows(QModelIndex(), position, position + rows - 1)
        self.flows = (self.flows[:position] +
                      self.flows[position + rows:])
        self.endRemoveRows()
        self.dirty = True
        return True


    def load(self):
        exception = None
        fh = None
        try:
            if self.filename.isEmpty():
                raise IOError, "no filename specified for loading"
            fh = QFile(self.filename)
            if not fh.open(QIODevice.ReadOnly):
                raise IOError, unicode(fh.errorString())
            stream = QDataStream(fh)
            magic = stream.readInt32()
            if magic != MAGIC_NUMBER:
                raise IOError, "unrecognized file type"
            fileVersion = stream.readInt16()
            if fileVersion != FILE_VERSION:
                raise IOError, "unrecognized file type version"
            self.flows = []
            while not stream.atEnd():
                opendpi = QString()
                l4 = QString()
                processname = QString()
                fivetuple = QString()
                features = QString()
                payload = QString()
                stream >> opendpi >> l4 >> processname >> fivetuple >> features >> payload
                #features = stream.readInt32()
                self.flows.append(Flow(opendpi, l4, processname,fivetuple, features, payload))
                self.l4set.add(unicode(l4))
                self.processnameset.add(unicode(processname))
                self.featuresset.add(features)
                self.payloadfilenameset.add(payload)
            self.dirty = False
        except IOError, e:
            exception = e
        finally:
            if fh is not None:
                fh.close()
            if exception is not None:
                raise exception


    def save(self):
        exception = None
        fh = None
        try:
            if self.filename.isEmpty():
                raise IOError, "no filename specified for saving"
            fh = QFile(self.filename)
            if not fh.open(QIODevice.WriteOnly):
                raise IOError, unicode(fh.errorString())
            stream = QDataStream(fh)
            stream.writeInt32(MAGIC_NUMBER)
            stream.writeInt16(FILE_VERSION)
            stream.setVersion(QDataStream.Qt_4_1)
            for flow in self.flows:
                stream << flow.opendpi << flow.l4 << flow.processname << flow.fivetuple << flow.features <<flow.payloadfilename
            self.dirty = False
        except IOError, e:
            exception = e
        finally:
            if fh is not None:
                fh.close()
            if exception is not None:
                raise exception


class FlowDelegate(QStyledItemDelegate):

    def __init__(self, parent=None):
        super(FlowDelegate, self).__init__(parent)


    def paint(self, painter, option, index):
        QStyledItemDelegate.paint(self, painter, option, index)


    def sizeHint(self, option, index):
        return QStyledItemDelegate.sizeHint(self, option, index)




def generateFakeFlows():
    for opendpi, l4, processname, fivetuple,  features, payloadfilename in (
( "unknown_pplive", "tcp"   , "pplive.exe"   ,"1.1.1.1-111-tcp-2.2.2.2-222", "a,b,c,d,e,f,g", "/home/flykeysky/dataset/1.fldata"),
( "unknown_pplive", "tcp"   , "pplive.exe"   ,"1.1.1.1-111-tcp-2.2.2.2-223", "a,b,c,d,e,f,g", "/home/flykeysky/dataset/2.fldata"),
( "unknown_pplive", "udp"   , "pplive.exe"   ,"1.1.1.1-111-udp-2.2.2.2-222",  "a,b,c,d,e,f,g", "/home/flykeysky/dataset/3.fldata"),
( "pplive"        , "tcp"   , "pplive.exe"   ,"1.1.1.1-111-tcp-2.2.2.2-224",  "a,b,c,d,e,f,g", "/home/flykeysky/dataset/4.fldata"),
( "pplive"        ,"udp"    ,  "thunder.exe" , "1.1.1.1-111-tcp-2.2.2.2-222", "a,b,c,d,e,f,g", "/home/flykeysky/dataset/5.fldata"),
( "unknown"       , "tcp"   , "thunder.exe"  , "1.1.1.1-111-tcp-2.2.2.2-222",  "a,b,c,d,e,f,g", "/home/flykeysky/dataset/6.fldata"),
( "unknown_thunder", "tcp"  ,"thunder.exe"   , "1.1.1.1-111-tcp-2.2.2.2-222",  "a,b,c,d,e,f,g", "/home/flykeysky/dataset/7.fldata"),
): yield Flow(opendpi, l4, processname, fivetuple, features, payloadfilename)

